import os
import matplotlib
from matplotlib import pyplot as PLT
from datetime import datetime
from scipy.stats import linregress

def Polyfnc(x, m, b):
        return m*x + b
    
def PlotParameters(database, today):
    """ Plots each gaussian fit parameters vs date
         p[0], p[1], p[2] = amplitude , center_x, sigma: for peak A
         p[3], p[4] = amplitude, sigma: for peak B
         p[5] = polynomial offset
         p[6], p[7], p[8] = ax^2 + bx + c """
    
    parameters = {}
    parameters['p0'] = []
    parameters['p1'] = []
    parameters['p2'] = []
    parameters['p3'] = []
    parameters['p4'] = []
    parameters['p5'] = []
    parameters['p6'] = []
    parameters['p7'] = []
    parameters['p8'] = []
    filelist = []
    date = []
    intdate = []
        
    # get parameters for each run
    for run in database:
        if run.Direction() != 'in': continue
        if run.Source() != 'weakTh-228': continue
        if run.Station() != 'S5': continue
        if run.NewParms() == None: continue
        
        np = run.NewParms()
        
        parameters['p0'].append(np[0])
        parameters['p1'].append(np[1])
        parameters['p2'].append(np[2])
        parameters['p3'].append(np[3])
        parameters['p4'].append(np[4])
        parameters['p5'].append(np[5])
        parameters['p6'].append(np[6])
        parameters['p7'].append(np[7])
        parameters['p8'].append(np[8])
        
        filelist.append(run.Filename().split('.')[0].split(' ')[2].split('_')[0])
        rundate = datetime.strptime(str(run.Datetime()).split(' ')[0], '%Y-%m-%d')
        date.append(rundate)
    
    # convert date to int for plotting
    for d in date:
        intd = d - min(date)
        intdate.append(intd.days)
    
    # fit trendline to each parameter
    l0 = linregress(intdate, parameters['p0'])
    l1 = linregress(intdate, parameters['p1'])
    l2 = linregress(intdate, parameters['p2'])
    l3 = linregress(intdate, parameters['p3'])
    l4 = linregress(intdate, parameters['p4'])
    l5 = linregress(intdate, parameters['p5'])
    l6 = linregress(intdate, parameters['p6'])
    l7 = linregress(intdate, parameters['p7'])
    l8 = linregress(intdate, parameters['p8'])
    
    # Convert trendline equation to data array
    linearfits = {}
    linearfits['p0'] = []
    linearfits['p1'] = []
    linearfits['p2'] = []
    linearfits['p3'] = []
    linearfits['p4'] = []
    linearfits['p5'] = []
    linearfits['p6'] = []
    linearfits['p7'] = []
    linearfits['p8'] = []
    for thing in intdate:
        linearfits['p0'].append(Polyfnc(thing, l0[0], l0[1]))
        linearfits['p1'].append(Polyfnc(thing, l1[0], l1[1]))
        linearfits['p2'].append(Polyfnc(thing, l2[0], l2[1]))
        linearfits['p3'].append(Polyfnc(thing, l3[0], l3[1]))
        linearfits['p4'].append(Polyfnc(thing, l4[0], l4[1]))
        linearfits['p5'].append(Polyfnc(thing, l5[0], l5[1]))
        linearfits['p6'].append(Polyfnc(thing, l6[0], l6[1]))
        linearfits['p7'].append(Polyfnc(thing, l7[0], l7[1]))
        linearfits['p8'].append(Polyfnc(thing, l8[0], l8[1]))
    
    
    
    ## ____ Plots ____ ##
    box = [0.14, 0.14, 0.76, 0.76]
    
    markers = ['o', '^', 'D', 'o', 'D', 's', 'p', 'h', 'd'] 
    colors = ['b', 'g', 'r', 'r', 'g', 'b', 'r', 'g', 'b']
    ylabels = ['A6 Peak Amplitude (in.oz)', 'A6 Peak Location (cm)', 'Gaussian Width (cm)', 
    'A7 Peak Amplitude (in.oz)', 'Gaussian Width (cm)', 'Polynomial Offset (cm)', 'Leading Coefficient', 
    'First Power Coefficient', 'Constant Coefficient']
    titles = ['A6 Peak Amplitude Versus Date', 'A6 Peak Location Versus Date', 
    'Gaussian Width Versus Date for Peak at Bend A6', 'A7 Peak Amplitude (in.oz)',
    'Gaussian Width Versus Date for Peak at Bend A7', 'Polynomial Position Offset Versus Date', 
    'Second Order Polynomial Leading Coefficient Versus Date',
    "Second Order Polynomial's First Power Coefficient Versus Date", 
    "Second Order Polynomial's Constant Coefficient"]
    
    for p in parameters:
        print list(p)
        fig = PLT.figure(figsize = (15, 8), dpi = 150)
        ax = fig.add_axes(box)
        ax.set_ylabel(ylabels[int(list(p)[-1])])
        ax.set_xlabel('Date')
        PLT.title(titles[int(list(p)[-1])])
        ax.grid()
        ax.scatter(date, parameters[p], color = colors[int(list(p)[-1])], marker = markers[int(list(p)[-1])])
        ax.plot(date, linearfits[p], '-', color = colors[int(list(p)[-1])])
        fig.savefig('averagePlots/' + p + '_vs_date_' + today + '.png')
    
    
    
    fig = PLT.figure(figsize = (15, 8), dpi = 150)
    axa = fig.add_axes(box)
    axa.set_ylabel('Gaussian Width (cm)')
    axa.set_xlabel('Date')
    PLT.title('Gaussian Width Versus Date')
    axa.grid()
    axa.scatter(date, parameters['p2'], color = 'r', marker = 'D', label = 'Peak at Bend A6')
    axa.plot(date, linearfits['p2'], '-', color = 'r')
    axa.scatter(date, parameters['p4'], color = 'g', marker = 'D', label = 'Peak at Bend A7')
    axa.plot(date, linearfits['p4'], '-', color = 'g')
    h, l = axa.get_legend_handles_labels()
    #for i, txt in enumerate(filelist):
        #axa.annotate(txt, (date[i], parameters['p4'][i]), position = (date[i], parameters['p4'][i]), size = 'x-small', alpha = 0.5, rotation = 45)
    PLT.legend(h, l, 'upper left')
    fig.savefig('averagePlots/p2_p4_vs_date_' + today + '.png')
    
    figb = PLT.figure(figsize = (15, 8), dpi = 150)
    axb = figb.add_axes(box)
    axb.set_ylabel('Amplitude (in.oz)')
    axb.set_xlabel('Date')
    PLT.title('Amplitude of Peaks at A6 and A7 Versus Date')
    axa.grid()
    axb.scatter(date, parameters['p0'], color = 'b', marker = 'o', label = 'A6 Peak')
    axb.plot(date, linearfits['p0'], '-', color = 'b')
    axb.scatter(date, parameters['p3'], color = 'r', marker = 'o', label = 'A7 Peak')
    axb.plot(date, linearfits['p3'], '-', color = 'r')
    h, l = axb.get_legend_handles_labels()
    PLT.legend(h, l, 'upper left')
    figb.savefig('averagePlots/p0_p3_vs_date_' + today + '.png')
    
    
    